from torch.utils.data import Dataset
import torch

class MotoAngleSolverDataset(Dataset):
    
    def __init__(self, path):
        super(Dataset, self).__init__()
        with open(path, 'r') as file:
            self.lines = [line.strip() for line in file]
        datas = [line.split() for line in self.lines]
        self.x = [data[-16:] for data in datas]
        self.y = [data[:6] for data in datas]
        # 将字符串数据转化为浮点数并转换为 Tensor
        self.x = torch.tensor([list(map(float, x)) for x in self.x], dtype=torch.float32)
        self.y = torch.tensor([list(map(float, y)) for y in self.y], dtype=torch.float32)

    def __getitem__(self, idx):
        return self.x[idx], self.y[idx]

    def __len__(self,):
        return len(self.lines)


if __name__ == "__main__":
    path = "./dataset/train_data.txt"
    data = MotoAngleSolverDataset(path)
    x, y = data[0]
    print(x)
    print(y)
    print(len(data))
    print("-"*80)
